# comp_state_pal <- c("#247C90", "red")
# 
# names(comp_state_pal) <- c('Probabilistic Bias Analysis', "Covidestim")
# 
# 
# vlab_for_pal <- unique(subset$vlabel) %>% as.character()

color_for_bias <- pal[[which(names(pal)=="$Priors\\,Do\\,Not\\,Vary\\,by\\,State\\,and\\,Date$")]]

comp_state_pal<- c(color_for_bias,"red")
names(comp_state_pal) <- c('Probabilistic Bias Analysis', "Covidestim")


subset %>%
   # mutate(vlabel = gsub("$\\,$CTIS Priors that Do Not Vary by State or Date$\\,$",
   #                     "$Priors\\,Do\\,Not\\,Vary\\,by\\,State\\,and\\,Date$", as.character(vlabel))) %>% 
   ggplot() +
  geom_ribbon(aes(x = date, 
             ymin = exp_cases_lb,
             ymax = exp_cases_ub,
             fill = 'Probabilistic Bias Analysis'),
             alpha = .8,
             show.legend=FALSE) +
  geom_ribbon(aes(x = date,
                  fill = 'Covidestim',
                  ymin = infections.lo,
                  ymax = infections.hi),
              alpha = .6) +
  geom_vline(xintercept=end_old_model,
             linetype=2,
             color="darkred",
             linewidth=.5) +
  facet_wrap(~name, ncol=4, scales = "free",
              labeller= labeller(vlabel =as_labeller(TeX, default=label_parsed))) +
  scale_y_continuous(labels = comma) +
  scale_x_date(date_breaks = "3.5 months", 
               date_labels = "%b %Y") +
  theme_c(axis.text.x=element_text(size=7),
          axis.text.y = element_text(size=4.5),
          legend.position = "top",
          legend.text=element_text(size=12),
          strip.text.x=element_text(size=11,
                                    face='plain',
                                    color='white')) +
  scale_fill_manual(values =comp_state_pal, name='')  +
  labs(y = "Estimated Infections",
       title = paste0("Estimated Infections by State")) +
  guides(color = guide_legend(override.aes = list(linewidth =2),
                              nrow=2,
                              byrow=TRUE))

ggsave(here('figures/intervals-and-covidestim-all-states.pdf'))

# states_with_lag

subset %>%
    filter(state %in% states_with_lag & biweek <= 25) %>%
   ggplot() +
  geom_ribbon(aes(x = date, 
             ymin = exp_cases_lb,
             ymax = exp_cases_ub,
             fill = 'Probabilistic Bias Analysis'),
             alpha = .8,
             show.legend=FALSE) +
  geom_ribbon(aes(x = date,
                  fill = 'Covidestim',
                  ymin = infections.lo,
                  ymax = infections.hi),
              alpha = .6) +
  # geom_vline(xintercept=end_old_model,
  #            linetype=2,
  #            color="darkred",
  #            linewidth=.5) +
  facet_wrap(~name, ncol=3, scales = "free",
              labeller= labeller(vlabel =as_labeller(TeX, default=label_parsed))) +
  scale_y_continuous(labels = comma) +
  scale_x_date(date_breaks = "3.5 months", 
               date_labels = "%b %Y") +
  theme_c(axis.text.x=element_text(size=7),
          axis.text.y = element_text(size=4.5),
          legend.position = "top",
          legend.text=element_text(size=12),
          strip.text.x=element_text(size=11,
                                    face='plain',
                                    color='white')) +
  scale_fill_manual(values =comp_state_pal, name='')  +
  labs(y = "Estimated Infections",
       title = paste0("Estimated Infections by State\nin States with a Lag")) +
  guides(color = guide_legend(override.aes = list(linewidth =2),
                              nrow=2,
                              byrow=TRUE))

ggsave(here('figures/intervals-and-covidestim-states-lag.pdf'))

############################################
# POSITIVE TESTS AND TEST POSITIVITY
############################################
pltlist <- subset %>%
  filter(biweek<=25) %>%
  filter(state %in% states_with_lag) %>%
    mutate(testpos = positive/total) %>%
  group_by(across(-date)) %>%
  summarize(date=min(date) + days(7)) %>%
  group_by(state) %>%
  mutate(adj = (1/max(testpos)) * max(positive)) %>%
  group_split() %>%
  map(~ {
    adj <-unique(.x$adj)
    .x %>%
      ggplot(aes(x=date))+
    geom_line(aes(y=positive, color = 'Positive Tests')) +
      geom_point(aes(y=positive, color = 'Positive Tests'),
              alpha=.5,
              size=.8) +
    geom_line(aes(y=testpos*adj, color='Test Positivity'),
              linetype=2) +
      geom_point(aes(y=testpos*adj, color='Test Positivity'),
              alpha=.5,
              size=.8) +
    #scale_linetype_discrete(breaks = c(1,1,2)) +
    scale_color_manual(values=c(
                                "Positive Tests"= "#279143",
                                "Test Positivity"= '#E38E4F')) +
    guides(color = guide_legend(override.aes = list(linetype = c(1,2),
                                                    linewidth=c(2,2)))) +
    theme_c() +
    theme(axis.title = element_text(size=10),
          axis.text =element_text(size=9),
          legend.position="none") +
    scale_y_continuous(labels=comma, 
                       sec.axis = sec_axis(
                         ~./adj,
                         name = 'Test Positivity')) +
    labs(color ='',
         y="Positive Tests",
         x= "Date") +
      facet_wrap(~name)
        
  })

legend <- get_legend(
  pltlist[[1]] +
    theme(legend.position = "top",
          legend.text=element_text(size=16))
)


title_gg <- ggplot() + 
  labs(title = "Comparing Trends in Positive Tests and Test Positivity") + 
  theme(plot.title=element_text(face="bold",
                                hjust = .5, 
                                size = 18,
                                margin =margin(5,0,2,0)))


plt <- cowplot::plot_grid(plotlist=pltlist, ncol=3) 

plot_grid(title_gg, legend, plt, rel_heights=c(.03, .03,.94), ncol=1)

ggsave(here('figures/testpos-lag.pdf'))

subset %>% 
   ggplot() +
  geom_ribbon(aes(x = date, 
             ymin = exp_cases_lb/positive,
             ymax = exp_cases_ub/positive,
             fill = vlabel),
             alpha = .9,
             show.legend=FALSE,
             fill= "#247C90") +
  facet_wrap(~state, ncol=5, 
              labeller= labeller(vlabel =as_labeller(TeX, default=label_parsed))) +
  scale_y_continuous(labels = comma) +
  scale_x_date(date_breaks = "3 months", 
               date_labels = "%b %Y") +
  theme_c(axis.text.x=element_text(size=6),
          axis.text.y = element_text(size=8),
          legend.position = "top") +
  scale_fill_manual(values =pal,
                    labels = TeX(names(pal)),
                    name='',
                    breaks="Covidestim")  +
  labs(y = "Ratio of Estimated Infections to Observed",
       x="",
       title = paste0("Ratio of Estimated Infections to Observed by State")) +
  guides(color = guide_legend(override.aes = list(linewidth =2),
                              nrow=2,
                              byrow=TRUE))

ggsave(here('figures/ratio-estimated-to-observed-all-states.pdf'))

max_val <- subset %>% 
  # left_join(codes) %>%
  filter(biweek %in% c(7, 13,27)) %>%
  mutate(ratio=exp_cases_ub/positive) %>%
  pull(ratio) %>%
  max()
  
  
plt1 <- subset %>% 
  # left_join(codes) %>%
  filter(biweek == 7) %>%
  ggplot(aes(y =fct_reorder(name,exp_cases_median/positive),
             xmin = exp_cases_lb/positive,
             xmax = exp_cases_ub/positive)) +
  geom_errorbar() +
  # geom_point(aes(y=state_name,x=exp_cases_median/positive),
  #            color = "darkred", size=.8, alpha=.5) +
  theme_c() + 
  theme(axis.text.y=element_text(size=10),
        axis.text.x = element_text(size=11),
        axis.title = element_text(size = 17),
        plot.subtitle = element_text(size=14),
        plot.title = element_text(size=17, face="plain")) +
  labs(y="", x="Ratio of Estimated Infections\nto Observed",
       title ="During Two-week Period\nin the Alpha Wave",
       subtitle = "March 26, 2021 through April 8, 2021") +
  xlim(0,max_val+2)

plt1

subset %>%
  mutate(ratio = exp_cases_median/positive) %>%
  group_by(biweek) %>%
  summarize(mean = mean(ratio),
            mindate=min(date),
            maxdate=max(date)) %>%
  arrange(desc(mean)) 

plt2 <- subset %>% 
  # left_join(codes) %>%
  filter(biweek == 13) %>%
  ggplot(aes(y =fct_reorder(name,exp_cases_median/positive),
             xmin = exp_cases_lb/positive,
             xmax = exp_cases_ub/positive)) +
  geom_errorbar() +
  # geom_point(aes(y=state_name,x=exp_cases_median/positive),
  #            color = "darkred", size=.8, alpha=.5) +
  theme_c() + 
 theme(axis.text.y=element_text(size=10),
        axis.text.x = element_text(size=11),
        axis.title = element_text(size = 17),
        plot.subtitle = element_text(size=14),
        plot.title = element_text(size=17, face="plain")) +
  labs(y="", x="Ratio of Estimated Infections\nto Observed",
       title ="During Two-week Period\nin the Delta Wave",
       subtitle = "June 18, 2021 through July 1, 2021") +
  xlim(0,max_val+2)

plt2

subset %>% 
  left_join(codes) %>%
  filter(biweek == 27) %>%
  pull(date) %>% 
  range()


subset %>% 
  left_join(codes) %>%
  filter(biweek == 13) %>%
  pull(date) %>% 
  range()


plt3 <- subset %>% 
  # left_join(codes) %>%
  filter(biweek == 27) %>%
  ggplot(aes(y =fct_reorder(name,exp_cases_median/positive),
             xmin = exp_cases_lb/positive,
             xmax = exp_cases_ub/positive)) +
  geom_errorbar() +
  theme_c() + 
  theme(axis.text.y=element_text(size=10),
        axis.text.x = element_text(size=11),
        axis.title = element_text(size = 17),
        plot.subtitle = element_text(size=14),
        plot.title = element_text(size=17, face="plain")) +
  labs(y="", x="Ratio of Estimated Infections\nto Observed",
       title ="During Two-week Period\nin the Omicron Wave",
       subtitle = "December 31, 2021 through January 14, 2022") +
  xlim(0,max_val+2)

plt3

plt <- plot_grid(plt1,plt2, plt3, nrow=1, align="none", labels="auto", label_size=19)

title <- ggplot() +
  labs(title = "Ratio of Estimated Infections to Observed Infections") +
  theme_c(plot.title=element_text(size=25))

plot_grid(title, plt, nrow=2, 
          rel_heights= c(.05,.95))

ggsave(here('figures/ratio-estimated-to-observed-multiple-waves.pdf'))

cat("Time interval with the highest ratio of estimated cases to observed: ")

subset %>%
  group_by(biweek) %>%
  mutate(ratio=exp_cases_median/positive) %>%
  summarize(mean_ratio=mean(ratio),
         date = paste(range(date),collapse=" to ")) %>%
  slice_max(n=1, order_by=mean_ratio)

subset %>% 
  # left_join(codes) %>%
  filter(biweek %in% c(7, 13, 27))  %>%
  select(-date) %>%
  distinct() %>%
  group_by(state) %>%
  mutate(ord = empirical_testpos[which(biweek==7)]) %>%
  mutate(biweek=factor(biweek)) %>%
  ggplot(aes(y=fct_reorder(state,ord),
             x=empirical_testpos,fill=biweek)) +
  geom_bar(stat="identity",position="dodge")+
  theme_c() +
  theme_c() +
  labs(title = "Testing Positivity by State",
       y="",
       x="Biweekly Testing Positivity",
       fill="") +
  scale_x_continuous(n.breaks=10)

subset %>% 
  # left_join(codes) %>%
  filter(biweek %in% c(7, 13, 27))  %>%
  group_by(biweek) %>%
  mutate(daterange= paste(format(min(date), "%B %d %Y"),
                          "to", 
                          format(max(date), "%B %d %Y")),
         daterange=factor(
           daterange,
           levels=c(
             "March 26 2021 to April 08 2021",
             "June 18 2021 to July 01 2021",
             "December 31 2021 to January 14 2022"))) %>%
  select(-date) %>%
  distinct() %>%
  group_by(state) %>%
  mutate(testrate=total/population,
         ord = testrate[which(biweek==7)]) %>%
  mutate(biweek=factor(biweek)) %>%
  ggplot(aes(y=fct_reorder(name,ord,.desc=TRUE),
             x=testrate,fill=daterange)) +
  geom_bar(stat="identity",position="dodge") +
  theme_c() +
  labs(title = "Testing Rate by State",
       y="",
       x="Biweekly Testing Rate",
       fill="") +
  scale_x_continuous(n.breaks=10)

compare_testrates <- subset %>%
  mutate(testrate=total/population) %>%
  select(-date) %>%
  distinct() %>%
  filter(biweek %in% c(7, 13, 27))  %>%
  select(biweek,state,testrate) %>%
  pivot_wider(names_from=biweek, values_from=testrate, names_prefix="biweek_") %>%
  mutate(omicron_over_alpha = biweek_27/biweek_7,
         omicron_over_delta = biweek_27/biweek_13) %>%
  select(state,contains("omicron")) %>%
  pivot_longer(c(omicron_over_alpha,omicron_over_delta)) %>%
  group_by(name) %>%
  summarize(mean_val = mean(value))

cat("The testing rate during this two-week interval during the omicron wave was", compare_testrates %>% filter(name=="omicron_over_alpha") %>% pull(mean_val), "times the testing rate during this time interval in the alpha wave.") 

cat("The testing rate during this two-week interval during the omicron wave was", compare_testrates %>% filter(name=="omicron_over_delta") %>% pull(mean_val), "times the testing rate during this time interval in the delta wave.") 

state_population_link <- "https://www2.census.gov/programs-surveys/popest/datasets/2010-2019/state/detail/SCPRC-EST2019-18+POP-RES.csv"
state_pop <- read_csv(state_population_link)
statecodes <- read_csv(here("data/demographic/statecodes.csv"))
  
state_pop <- state_pop %>%
    left_join(statecodes, by = c("NAME" = "state")) %>%
    select(population = POPESTIMATE2019,
           state = code,
           name = NAME) %>%
    filter(!is.na(state))
  
state_deaths <- tar_read(state_deaths, store=targets_dir)


set.seed(123)
state_sample <- sample(unique(state_deaths$state), 5)

state_deaths %>%
#  filter(state %in% state_sample) %>%
  ggplot(aes(x=date,y=deaths/cases)) +
  geom_line() +
  facet_wrap(~state) +
  theme_c() +
  geom_vline(xintercept=end_old_model, color='darkred', linetype=2)

comp_deaths <- state_deaths %>%
  mutate(omicron=ifelse(date <= mdy("12-02-2021"),
                        "pre_omicron", "omicron")) %>%
  group_by(state, omicron) %>%
  summarize(cumulative_deaths = sum(deaths),
            cumulative_cases= sum(cases)) %>%
  pivot_wider(names_from = omicron, 
              values_from = c(cumulative_deaths, cumulative_cases)) %>%
  mutate(pre_deaths_over_cases = cumulative_deaths_pre_omicron/cumulative_cases_pre_omicron,
         omicron_deaths_over_cases = cumulative_deaths_omicron/cumulative_cases_omicron) 



comp_deaths %>%
  select(state,pre_deaths_over_cases,omicron_deaths_over_cases) %>%
  pivot_longer(contains("over_cases")) %>%
  group_by(state) %>%
  mutate(pre = value[which(name=="pre_deaths_over_cases")]) %>%
  mutate(name = case_when(
    grepl("omicron", name) ~ "Omicron",
    grepl("pre", name) ~ "Before Omicron"
  )) %>%
  ggplot(aes(y=fct_reorder(state,pre),  x= value, fill=name)) +
  geom_bar(stat='identity', position='dodge') +
  theme_c(legend.title = element_text(face="bold", size=14)) +
  labs(x = TeX("$\\frac{Total\\,Deaths}{Total\\,Cases}$"),
       y="",
       fill = "Time Period",
       title = "Total Deaths Over Total Cases in the Time Period Before Omicron\n Compared to During Omicron") +
  scale_x_continuous(n.breaks=7)

comp_deaths %>%
  select(state,pre_deaths_over_cases,omicron_deaths_over_cases) %>%
  mutate(pct_change =( omicron_deaths_over_cases -pre_deaths_over_cases ) / pre_deaths_over_cases ) %>% 
  ggplot(aes(y=fct_reorder(state,pct_change),  x= pct_change)) +
  geom_bar(stat='identity', position='dodge') +
  theme_c(legend.title = element_text(face="bold", size=11),
          axis.text.x=element_text(size=14),
          plot.subtitle = element_text(size=16,
                                       face="plain", 
                                       margin=margin(0,0,4,0))) +
  labs(x = TeX("$\\frac{Total\\,Deaths}{Total\\,Cases}$"),
       y="",
       fill = "Time Period",
       title = TeX("Percent Change in $\\frac{Total\\,Deaths}{Total\\,Cases}$ from the Time Period Before Omicron"),
       subtitle="Compared to During Omicron") +
  scale_x_continuous(labels=scales::percent)

comp_deaths %>%
  select(state,pre_deaths_over_cases,omicron_deaths_over_cases) %>%
  mutate(pct_change =( omicron_deaths_over_cases - pre_deaths_over_cases ) / pre_deaths_over_cases ) %>%
  left_join(state_overlaps, by = c('state' = 'name')) %>%
  group_by(state) %>%
  summarize(mean_overlap=mean(fraction_overlap),
            pct_change=unique(pct_change)) %>%
  ggplot(aes(x=pct_change,y=mean_overlap)) +
  geom_point(size=2) +
  theme_c(  plot.subtitle = element_text(size=16,
                                       face="plain", 
                                       margin=margin(0,0,4,0))) +
  scale_x_continuous(labels=scales::percent) +
  labs(x="Percent Change", y = "Mean Overlap\n(By State)", 
        title = TeX("Relationship Between the Percent Change in $\\frac{Total\\,Deaths}{Total\\,Cases}$ from the Time Period Before Omicron"),
       subtitle="Compared to During Omicron versus the Mean Overlap")

comp_deaths %>%
  rename(name=state) %>%
  ungroup() %>%
  left_join(state_pop[,c("name", "population")]) %>%
  mutate(frac_infected = cumulative_cases_pre_omicron/population) %>%
  left_join(state_overlaps[,c("name","fraction_overlap", "biweek")]) %>%
  filter(biweek >= 25) %>%
  group_by(name) %>%
  summarize(mean_overlap=mean(fraction_overlap),
            frac_infected=unique(frac_infected)) %>%
  ggplot(aes(x=frac_infected, y =mean_overlap)) +
  geom_point() +
  theme_c()

covidestim_link <- "https://covidestim.s3.us-east-2.amazonaws.com/latest/state/estimates.csv"
  
covidestim <- read_csv(covidestim_link)
  
# join data from each source to include dates before and after 2021-12-02
covidestim <- covidestim %>%
    select(date, contains("infections"), state) %>%
    filter(year(date) > 2020) %>%
    mutate(week = week(date), year = year(date))


covidestim %>%
  group_by(state) %>%
  summarize(infections=sum(infections))%>%
  rename(name=state) %>%
  left_join(state_pop[,c("name", "population")]) %>%
  mutate(frac_infected = infections/population) %>%
  left_join(state_overlaps[,c("name","fraction_overlap", "biweek")]) %>%
  filter(biweek >= 25) %>%
  group_by(name) %>%
  summarize(mean_overlap=mean(fraction_overlap),
            frac_infected=unique(frac_infected)) %>%
  ggplot(aes(x=frac_infected, y =mean_overlap)) +
  geom_point() +
  theme_c()